//===- PiNodeInsertion.cpp - Insert Pi nodes into a program ---------------===//
// 
//                     The LLVM Compiler Infrastructure
//
// This file was developed by the LLVM research group and is distributed under
// the University of Illinois Open Source License. See LICENSE.TXT for details.
// 
//===----------------------------------------------------------------------===//
//
// PiNodeInsertion - This pass inserts single entry Phi nodes into basic blocks
// that are preceded by a conditional branch, where the branch gives
// information about the operands of the condition.  For example, this C code:
//   if (x == 0) { ... = x + 4;
// becomes:
//   if (x == 0) {
//     x2 = phi(x);    // Node that can hold data flow information about X
//     ... = x2 + 4;
//
// Since the direction of the condition branch gives information about X itself
// (whether or not it is zero), some passes (like value numbering or ABCD) can
// use the inserted Phi/Pi nodes as a place to attach information, in this case
// saying that X has a value of 0 in this scope.  The power of this analysis
// information is that "in the scope" translates to "for all uses of x2".
//
// This special form of Phi node is referred to as a Pi node, following the
// terminology defined in the "Array Bounds Checks on Demand" paper.
//
// As a really trivial example of what the Pi nodes are good for, this pass
// replaces values compared for equality with direct constants with the constant
// itself in the branch it's equal to the constant.  In the case above, it would
// change the body to be "... = 0 + 4;"  Real value numbering can do much more.
//
//===----------------------------------------------------------------------===//
//
//Modified by Xiao Yu For llvm release 3.0
//

#define DEBUG_TYPE "pinode"

#include <list>
#include <string>
#include <sstream>

#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/InstrTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Analysis/DominanceFrontier.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Transforms/Utils/PromoteMemToReg.h"
#include "llvm/Support/raw_ostream.h"

namespace llvm {

STATISTIC(NumPINodeInserted, "Number of Pi nodes inserted");

namespace {

struct PiNodeInserter : public FunctionPass {

  static char ID;

  DominatorTree* DT;
  DominanceFrontier* DF;

  PiNodeInserter() : FunctionPass(ID), DT(NULL), DF(NULL) {}

  virtual bool doInitialization(Module &M) {
    NumPINodeInserted = 0;
    return false;
  }

  virtual bool runOnFunction(Function &F);

  virtual bool doFinalization(Module &M) {
    errs() << "[PINodeInsert]>>>Total number of PI/PHI nodes added: "
        << NumPINodeInserted << '\n';
    return false;
  }

  virtual void getAnalysisUsage(AnalysisUsage &AU) const {
    AU.addRequired<DominatorTree>();
    AU.addRequired<DominanceFrontier>();
  }

 private:
  // insertPiNode - Insert a Pi node/block for V. 
  PHINode* insertPiNode(unsigned Idx, Value* V,
                        BasicBlock* Pred, BasicBlock* Succ);

  // propagatePiNode - Update the Graph for the newly added Pi node
  void propagatePiNode(PHINode* Pi, Value* V,
                       BasicBlock* Pred, BasicBlock* Succ);

  //getReachingIdxFor - Get indices of edges that definitions in BB can reach
  std::set<unsigned int> getReachingIdxFor(PHINode* Phi, BasicBlock* BB);

  //replaceValueFor - Replace a use of V with New if New dominates the use
  void replaceValueIfDom(Value* V, Value* New, BasicBlock* BB);

  //insertNewPhiFor 
  //- Insert a new phi node at block FB for Phi in the place of V
  PHINode* insertNewPhiFor(PHINode* Phi, Value* V, BasicBlock* FB);
  
  //update 
  //- Update the newly inserted node Phi in place of V
  void update(PHINode* Phi, Value* V);

};

char PiNodeInserter::ID = 0;
static RegisterPass<PiNodeInserter> X("pinode", "Pi Node Insertion");
}

bool PiNodeInserter::runOnFunction(Function &F) {
  bool Changed = false;

  //Dominance information
  DT = &getAnalysis<DominatorTree>();
  DF = &getAnalysis<DominanceFrontier>();

  //Branches
  std::list<BranchInst*> CondBrList;
  std::list<SwitchInst*> SwitchList;
  for (Function::iterator B = F.begin(), E = F.end(); B != E; ++B) {
    TerminatorInst* TI = B->getTerminator();
    if (BranchInst* BI = dyn_cast<BranchInst>(TI))
      if (BI->isConditional())
        CondBrList.push_back(BI);
    if (SwitchInst* SI = dyn_cast<SwitchInst>(TI))
      SwitchList.push_back(SI);
  }


  //Conditional Branches
  for (std::list<BranchInst*>::iterator I = CondBrList.begin(), 
       E = CondBrList.end(); I != E; ++I) {
    BranchInst* BI = *I;
    // CmpInst
    if (CmpInst *CI = dyn_cast<CmpInst>(BI->getCondition())) {
      if (CI->getNumOperands() < 2)
        continue;
      if (!CI->isIntPredicate())
        continue;

      BasicBlock* TB = BI->getSuccessor(0);  // True block
      BasicBlock* FB = BI->getSuccessor(1);  // False block

      // Insert the Pi nodes for the first operand to the comparison...
      PHINode* Pi;
      Pi = insertPiNode(0, CI->getOperand(0), BI->getParent(), TB);
      Changed |= (Pi != NULL);
      propagatePiNode(Pi, CI->getOperand(0), BI->getParent(), TB);

      Pi = insertPiNode(1, CI->getOperand(0), BI->getParent(), FB);
      Changed |= (Pi != NULL);
      propagatePiNode(Pi, CI->getOperand(0), BI->getParent(), FB);

      // Insert the Pi nodes for the second operand to the comparison...
      Pi = insertPiNode(0, CI->getOperand(1), BI->getParent(), TB);
      Changed |= (Pi != NULL);
      propagatePiNode(Pi, CI->getOperand(1), BI->getParent(), TB);

      Pi = insertPiNode(1, CI->getOperand(1), BI->getParent(), FB);
      Changed |= (Pi != NULL);
      propagatePiNode(Pi, CI->getOperand(1), BI->getParent(), FB);
    }
  }

  //Switches
  for (std::list<SwitchInst*>::iterator I = SwitchList.begin(),
       E = SwitchList.end(); I != E; ++I) {
    SwitchInst* SI = *I;
    for (unsigned int i = 0; i < SI->getNumCases(); ++i) {
      BasicBlock* Succ = SI->getSuccessor(i);
      PHINode* Pi;
      Pi = insertPiNode(i, SI->getCondition(), SI->getParent(), Succ);
      Changed |= (Pi != NULL);
      propagatePiNode(Pi, SI->getCondition(), SI->getParent(), Succ);
    }
  }

  return Changed;
}


// getPhiNodeFor - Return the phi node if there is already one in BB for V.
static PHINode* getPhiNodeFor(Value *V, BasicBlock *BB) {
  for (Value::use_iterator I = V->use_begin(), E = V->use_end(); I != E; ++I)
    if (PHINode *PN = dyn_cast<PHINode>(*I))
      if (PN->getParent() == BB)
        return PN;
  return NULL;
}


// insertPiBlock
// Insert the pi node and return it if newly created, otherwise return NULL
PHINode* PiNodeInserter::insertPiNode(unsigned Idx, Value *V,
                                      BasicBlock* Pred, BasicBlock *Succ) {
  // Do not insert Pi nodes for constants!
  if (isa<Constant>(V)) return NULL;

  //Do not insert for pointer type
  if (V->getType()->getTypeID() == Type::PointerTyID) return NULL;

  //Check if a new block is needed
  BasicBlock* PiBB = Succ;
  TerminatorInst* TI = Pred->getTerminator();

  if (Succ->getSinglePredecessor() != Pred) {
    if (TI->getSuccessor(Idx) != Succ) {
      //new block already created
      PiBB = TI->getSuccessor(Idx);
    }
    else {
      //new BB before Succ
      PiBB = BasicBlock::Create(Pred->getParent()->getContext(), "ABCD_Pi_Block",
                                Pred->getParent(), Succ);

      //replace Pred with PiBB for Succ
      for (BasicBlock::iterator I = Succ->begin(), E = Succ->end();
           I != E; ++I) {
        if (PHINode* phi = dyn_cast<PHINode>(I)) {
          int idx = phi->getBasicBlockIndex(Pred);
//          while (idx != -1) {
//            phi->setIncomingBlock((unsigned)idx, PiBB);
//            idx = phi->getBasicBlockIndex(Pred);
//          } // end while
          phi->setIncomingBlock((unsigned)idx, PiBB);
        } // end if
      } // end for

      //branch from Pred to PiBB
      TI->setSuccessor(Idx, PiBB);

      //branch from PiBB to Succ
      BranchInst::Create(Succ, PiBB);
    } // end if not inserted
  }

  // Create a pi node in the pi block
  PHINode* PiNode = PHINode::Create(V->getType(), 1, V->getName() + "_Pi");
  PiNode->addIncoming(V, Pred);
  PiBB->getInstList().insert(PiBB->begin(), PiNode);

//  if (dyn_cast<SwitchInst>(TI)) {
//  errs() << "===============insertPiNode===========>" << '\n';
//  errs() << "Pred: " << '\n';
//  Pred->print(errs());
//  errs() << '\n';
//  errs() << "Succ: " << '\n';
//  Succ->print(errs());
//  errs() << '\n';
//  errs() << "PiBB: " << '\n';
//  PiBB->print(errs());
//  errs() << '\n';
//  errs() << "PiNode: " << '\n';
//  PiNode->print(errs());
//  errs() << '\n';
//  errs() << "==========================>" << '\n';
//  }

  ++NumPINodeInserted;
  return PiNode;
}

//getReachingIdxFor
//Return the indicies of edges that BB dominates
//Note that BB or Pred(Phi) might not in the dominance tree
std::set<unsigned int> 
PiNodeInserter::getReachingIdxFor(PHINode* Phi, BasicBlock* BB) {
  std::set<unsigned int> ret;
  for (unsigned int i = 0; i < Phi->getNumIncomingValues(); i++) {
    BasicBlock* Pred = Phi->getIncomingBlock(i);
    if (DT->getNode(Pred) && DT->getNode(BB)) {
      if (DT->dominates(BB, Pred)) {
        ret.insert(i);
        continue;
      }
    }

    //if Pred is BB
    if (Pred == BB) {
      ret.insert(i);
      continue;
    }

    //Pred not in DT, 
    //this implies Pred is a newly inserted Pi block
    //we can simply use its single predecessor here for the dominate
    //relationship
    if (!DT->getNode(Pred))
      Pred = Pred->getSinglePredecessor();

    //BB  not in DT
    //this implies BB is  a newly inserted Pi block
    //we can use its succ if it dominates its succ
    //otherwise, the definition in BB cannot reach Phi
    if (!DT->getNode(BB))
      if ((*succ_begin(BB))->getSinglePredecessor() == BB)
        BB = *succ_begin(BB);

    if (DT->dominates(BB, Pred))
      ret.insert(i);

  } // end for
  return ret;
}

inline std::string IntToString(int i) {
  std::stringstream ss;
  std::string ret;
  ss << i;
  ret = ss.str();
  return ret;
}

//find or insert a new version
inline int findOrInsertVersion(Value* V, std::map<Value*, int>& vm) {
  std::map<Value*, int>::iterator it = vm.find(V);
  int version = 0;
  if (it == vm.end()) {
    vm.insert(std::pair<Value*, int>(V, 0));
  }
  else {
    version = it->second;
    vm[V] ++;
  }
  return version;
}

//replaceValueIfDom()
//replace use of V with New under the dominance of BB
void PiNodeInserter::replaceValueIfDom(Value* V, Value* New, BasicBlock* BB) {
  std::vector<User*> Users(V->use_begin(), V->use_end());
  for (std::vector<User*>::iterator UI = Users.begin(), UE = Users.end(); 
       UI != UE; ++UI)
    if (Instruction *I = dyn_cast<Instruction>(*UI))
      if (DT->dominates(BB, I->getParent())) {
        if (dyn_cast<Instruction>(New->stripPointerCasts()) != I)
          I->replaceUsesOfWith(V, New);
      }
}

//insertNewPhiFor()
//insert a new phi in the frontier BB for Phi and V
//The Phi block might not in the dominance tree
//Return New if newly inserted
PHINode* PiNodeInserter::insertNewPhiFor(PHINode* Phi, Value* V, 
                                         BasicBlock* FB) {

  PHINode* New = getPhiNodeFor(V, FB);
  if (New) {
    //already a phi node there, replace the use of V
    //first get the blocks for New where Phi dominates
    std::set<unsigned int> idx_set = getReachingIdxFor(New, Phi->getParent());
    for (std::set<unsigned int>::const_iterator I = idx_set.begin(),
         E = idx_set.end(); I != E; ++I) {
      if (New->getIncomingValue(*I) == V) {
        New->setIncomingValue(*I, Phi);
      }
    }
    return NULL;
  } // if New != NULL

  //No New phi node here;
  
  //BasicBlock where V is defined
  BasicBlock* Vdef_BB;
  if (isa<Instruction>(V)) {
    Vdef_BB = (dyn_cast<Instruction>(V))->getParent();
  }
  else if (isa<Argument>(V)) {
    Vdef_BB = DT->getRoot();
  }
  else {
    errs() << "=====> Value Type not argument, not instruction" << '\n';
    return NULL;
  }

  //If this is a frontier of V and there is no phi
  //this implies there is no use of V, no need to add a new phi
  if (!DT->dominates(Vdef_BB, FB)) return NULL;

  //If there is a backedge and Vdef_BB happens to be the FB
  if (Vdef_BB == FB) return NULL;

  //Here a New phi should be inserted.
  static std::map<Value*, int> vm;
  int version = findOrInsertVersion(V, vm);
  unsigned num_pred = std::distance(pred_begin(FB), pred_end(FB));
  New = PHINode::Create(V->getType(), num_pred,
                        V->getName() + "_Phi." + IntToString(version),
                        FB->begin());
  //for the incoming edges where Phi dominates, incoming value is phi
  //otherwise, incoming value is V
  for(pred_iterator PI = pred_begin(FB), E = pred_end(FB); PI != E; ++PI) {
    New->addIncoming(V, *PI);
  }
  std::set<unsigned int> idx_set = getReachingIdxFor(New, Phi->getParent());
  for(unsigned int i = 0; i < num_pred; ++i) {
    if (idx_set.find(i) != idx_set.end())
      New->setIncomingValue(i, Phi);
  }

  return New;
}

//update()
//Replace the use of V with Phi or recursively insert new phi node
void PiNodeInserter::update(PHINode* Phi, Value* V) {
  //The block contains Phi should always in DT
  if (DT->getNode(Phi->getParent()) == NULL) {
    errs() << "Try to update a Pi node" << '\n';
  }

  //Replace
  replaceValueIfDom(V, Phi, Phi->getParent());

  //Insert
  DominanceFrontier::iterator it = DF->find(Phi->getParent());
  if (it == DF->end()) return;

  std::set<BasicBlock*> frontier = it->second;
  for (std::set<BasicBlock*>::const_iterator FI = frontier.begin(), 
       FE = frontier.end(); FI != FE; ++FI) {
    BasicBlock* FB = *FI;
    PHINode* New = insertNewPhiFor(Phi, V, FB);
    if (New == NULL) return;
    update(New, V);
  }
}

// propagatePiNode()
// Insert phi nodes for newly added pi block
// Replace uses of V with pi node
void PiNodeInserter::propagatePiNode(PHINode* Pi, Value* V,
                                     BasicBlock* Pred, BasicBlock* Succ) {
  if (Pi == NULL) return;

  //If Succ has multiple predecessor
  //Then a new Pi block will be added
  if (Succ->getSinglePredecessor() != Pred) {
    PHINode* New = insertNewPhiFor(Pi, V, Succ);
    if (New == NULL) return;
    update(New, V);
  }
  else {
    update(Pi, V);
  }
}

} // End llvm namespace
